#!/bin/bash 
export GLOBAL_LOG_LEVEL=3
export TF_CPP_MIN_LOG_LEVEL=3
export SLOG_PRINT_TO_STDOUT=0
execpath=$(dirname $(realpath $0))

IFS=',' read -ra NPUSARR <<<"${VISIBLE_IDS}"
echo "ADD---${NPUSARR[@]}"
export SOC_VERSION=Ascend910
export HCCL_CONNECT_TIMEOUT=600
export PRINT_MODEL=0
export RANK_SIZE=${#NPUSARR[@]}
export DEVICE_NUM=${#NPUSARR[@]}
export JOB_ID=$DLWS_JOB_ID
export RANK_START=0
export RANK_INDEX=0
export POD_NAME=${DLWS_JOB_ID}

logpath=~/work_dirs/distributed/$DLWS_JOB_ID/$DLWS_ROLE_IDX
mkdir -p $logpath
cd $logpath

if [ $DLWS_NUM_WORKER -gt 1 ]; then
  export POD_NAME=$DLWS_JOB_ID-worker-$DLWS_ROLE_IDX
  export SERVER_ID=${DLWS_ROLE_IDX}
  export RANK_START=$((DEVICE_NUM * SERVER_ID))
  export RANK_SIZE=$(($DLWS_NUM_WORKER * 8))
fi

export RANK_ID=$POD_NAME
ulimit -u unlimited
env

export AICPU_PROFILING_MODE=false
export PROFILING_OPTIONS=task_trace:training_trace
export FP_POINT=bert/embeddings/GatherV2
export BP_POINT=gradients/bert/embeddings/IdentityN_1_grad/UnsortedSegmentSum+
export RANK_TABLE_FILE=/home/$DLWS_USER_NAME/.npu/$DLWS_JOB_ID/hccl_tf.json
rm -rf output kernel_meta
INDEX=0
export RANK_INDEX=$INDEX
for device_phy_id in "${NPUSARR[@]}"; do
  export DEVICE_ID=$device_phy_id
  export RANK_DEVICE_INDEX=$INDEX
  INDEX=$((INDEX+1))

  python $execpath/src/pretrain/run_pretraining.py --bert_config_file=$execpath/configs/bert_base_config.json --max_seq_length=128 --max_predictions_per_seq=20 --train_batch_size=128 --learning_rate=1e-4 --num_warmup_steps=10000 --num_train_steps=500000 --optimizer_type=adam --manual_fp16=True --use_fp16_cls=True --input_files_dir=/data/dataset/storage/bookscorpus/training  --eval_files_dir=/data/dataset/storage/bookscorpus/test --npu_bert_debug=False --npu_bert_use_tdt=True --do_train=True --num_accumulation_steps=1 --npu_bert_job_start_file=None --iterations_per_loop=100 --save_checkpoints_steps=10000 --npu_bert_clip_by_global_norm=False --distributed=True --npu_bert_loss_scale=0 --output_dir=./output &
  echo "============ CMD ============"
  echo $CMD
  echo "============ CMD ============"
  eval $CMD
done
wait